import argparse
import os
import pickle

from utils import util


class BaseOptions:
    def __init__(self):
        super(BaseOptions, self).__init__()
        parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        parser = self.initialize(parser)
        # opt, parser = self.gather_options(parser)
        # self.opt = opt
        self.parser = parser

    def parse(self):
        self.print_options(self.opt)
        self.save_options(self.opt)
        return self.opt

    def gather_options(self, parser):
        # get the basic options
        opt, unknown = parser.parse_known_args()

        # if there is opt_file, load it.
        # The previous default options will be overwritten
        if opt.load_from_opt_file:
            parser = self.update_options_from_file(parser, opt)

        opt = parser.parse_args()
        return opt, parser

    def initialize(self, parser):
        # model config
        parser.add_argument('--name', type=str, default='sketch_colorization', help='name of the experiment. It '
                                                                                    'decides where to store samples '
                                                                                    'and models')
        parser.add_argument('--device', type=str, default='cuda')
        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
        parser.add_argument('--model', type=str, default='pix2pix', help='which model to use')


        parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
        parser.add_argument('--num_workers', type=int, default=8, help='thread of loading data')
        parser.add_argument('--use_tensorboard', type=bool, default=True, help='whether apply tensorboard to log')
        parser.add_argument('--load_from_opt_file', action='store_true')
        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')


        # Input or Output size
        parser.add_argument('--batch_size', type=int, default=1, help='batch size of model input')
        parser.add_argument('--load_size', type=int, default=256,
                            help='Scale images to this size. The final image will be converted to --image_size.')
        parser.add_argument('--image_size', type=int, default=256, help='output size of generated image')
        parser.add_argument('--preprocess_mode', type=str, default='scale_width_and_crop',
                            help='scaling and cropping of images at load time.', choices=(
            "resize_and_crop", "crop", "scale_width", "scale_width_and_crop", "scale_shortside",
            "scale_shortside_and_crop", "fixed", "none"))
        parser.add_argument('--aspect_ratio', type=float, default=1.0,
                            help='The ratio width/height. The final height of the load image will be crop_size/aspect_ratio')

        # dataset
        parser.add_argument('--dataroot', type=str, default='data/preprocessed_data')
        parser.add_argument('--dataset_mode', type=str, default='animepair')
        parser.add_argument('--no_pairing_check', action='store_true',
                            help='If specified, skip sanity check of correct label-image file pairing')

        # Augmentation
        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data argumentation')
        parser.add_argument('--maskmix', action='store_true', help='use mask in correspondence net')
        parser.add_argument('--mask_epoch', type=int, default=-1, help='useful when noise_for_mask is true, first train mask_epoch with mask, the rest epoch with noise')
        parser.add_argument('--noise_for_mask', action='store_true', help='replace mask with noise')
        parser.add_argument('--mask_noise', action='store_true', help='use noise with mask')

        parser.add_argument('--show_corr', action='store_true', help='if specified, save bilinear upsample correspondence')

        return parser

    def print_options(self, opt):
        name = opt.name + 'base options'
        message = ''
        message += '-------------- {0} ---------------\n'.format(name)
        for k, v in sorted(vars(opt).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '------------------------- End -------------------------'
        print(message)

    def option_file_path(self, opt, makedir=False):
        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
        if makedir:
            util.mkdirs(expr_dir)
        file_name = os.path.join(expr_dir, 'opt')
        return file_name

    def save_options(self, opt):
        file_name = self.option_file_path(opt, makedir=True)
        with open(file_name + '.txt', 'wt') as opt_file:
            for k, v in sorted(vars(opt).items()):
                comment = ''
                default = self.parser.get_default(k)
                if v != default:
                    comment = '\t[default: %s]' % str(default)
                opt_file.write('{:>25}: {:<30}{}\n'.format(str(k), str(v), comment))

        with open(file_name + '.pkl', 'wb') as opt_file:
            pickle.dump(opt, opt_file)

    def update_options_from_file(self, parser, opt):
        new_opt = self.load_options(opt)
        for k, v in sorted(vars(opt).items()):
            if hasattr(new_opt, k) and v != getattr(new_opt, k):
                new_val = getattr(new_opt, k)
                parser.set_defaults(**{k: new_val})
        return parser

    def load_options(self, opt):
        file_name = self.option_file_path(opt, makedir=False)
        new_opt = pickle.load(open(file_name + '.pkl', 'rb'))
        return new_opt

if __name__ == '__main__':
    base = BaseOptions()
    opt = base.parse()






